#include "WriteMemory.h"
#include "ExportFunc.h"

NTSTATUS WriteMemory(HANDLE pid, PVOID targetAddr, PVOID buf, SIZE_T size)
{
	if (!pid)
	{
		return STATUS_UNSUCCESSFUL;
	}
	if (targetAddr >= MmHighestUserAddress 
		|| ((ULONG64)targetAddr+size) >= MmHighestUserAddress
		|| buf >= MmHighestUserAddress
		|| ((ULONG64)targetAddr + size) < targetAddr
		)
	{
		return	STATUS_INVALID_PARAMETER_2;
	}
	PEPROCESS Process = {0};
	NTSTATUS status = PsLookupProcessByProcessId(pid, &Process);

	if (!NT_SUCCESS(status)) return STATUS_INVALID_PARAMETER_1;

	status = PsGetProcessExitStatus(Process);
	
	if (status != STATUS_PENDING) return STATUS_INVALID_PARAMETER_1;

	SIZE_T count = 0;
	PEPROCESS srcProcess = IoGetCurrentProcess();
	status = MmCopyVirtualMemory(srcProcess, buf, Process, targetAddr, size, UserMode, &count);

	if (NT_SUCCESS(status))
	{
		ObDereferenceObject(Process);
		return status;
	}
	KAPC_STATE kApcState = { 0 };
	
	KeStackAttachProcess(Process, &kApcState);

	ULONG64 baseAddress = (ULONG64)targetAddr;
	SIZE_T sizeLen = size;
	ULONG oldPro = 0;
	ULONG oldPro2 = 0;
  //status = NtProtectVirtualMemory(IoGetCurrentProcess(), &baseAddress, &sizeLen, PAGE_EXECUTE_READWRITE, &oldPro);                        
	status = NtProtectVirtualMemory(NtCurrentProcess(), &baseAddress, &sizeLen, PAGE_EXECUTE_READWRITE, &oldPro);

	if (NT_SUCCESS(status))
	{	
		status = MmCopyVirtualMemory(srcProcess, buf, Process, targetAddr, size, UserMode, &count);
		NtProtectVirtualMemory(NtCurrentProcess(), &baseAddress, &sizeLen, oldPro, &oldPro2);	
	}
	KeUnstackDetachProcess(&kApcState);
	if (!NT_SUCCESS(status))
	{
		ULONG64 value = wpoff();
		status = MmCopyVirtualMemory(srcProcess, buf, Process, targetAddr, size, UserMode, &count);
		wpon(value);
	}

	ObDereferenceObject(Process);
	return status;
}